package com.harmony.pilot.llm.provider;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.harmony.pilot.common.dto.LLMRequest;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * 智谱AI GLM-4 提供商实现
 * 支持GLM-4-6B等模型
 */
@Slf4j
@Component
public class GLM4Provider implements LLMProvider {
    
    private final OkHttpClient client;
    private final ObjectMapper objectMapper;
    
    @Value("${llm.glm4.api-key:}")
    private String apiKey;
    
    @Value("${llm.glm4.model:glm-4}")
    private String model;
    
    @Value("${llm.glm4.base-url:https://open.bigmodel.cn/api/paas/v4}")
    private String baseUrl;
    
    public GLM4Provider() {
        this.objectMapper = new ObjectMapper();
        this.client = new OkHttpClient.Builder()
                .connectTimeout(30, TimeUnit.SECONDS)
                .readTimeout(60, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .build();
    }
    
    @Override
    public String generate(LLMRequest request) {
        log.info("Generating with GLM-4");
        return callGLM4API(request, "generate");
    }
    
    @Override
    public String complete(LLMRequest request) {
        log.info("Completing with GLM-4");
        return callGLM4API(request, "complete");
    }
    
    @Override
    public String analyze(LLMRequest request) {
        log.info("Analyzing with GLM-4");
        return callGLM4API(request, "analyze");
    }
    
    @Override
    public String chat(LLMRequest request) {
        log.info("Chatting with GLM-4");
        return callGLM4API(request, "chat");
    }
    
    @Override
    public String getDefaultModel() {
        return model;
    }
    
    /**
     * 调用GLM-4 API
     */
    private String callGLM4API(LLMRequest request, String taskType) {
        try {
            // 构建请求体
            Map<String, Object> requestBody = buildRequestBody(request);
            String jsonBody = objectMapper.writeValueAsString(requestBody);
            
            log.debug("GLM-4 request: {}", jsonBody);
            
            // 构建HTTP请求
            Request httpRequest = new Request.Builder()
                    .url(baseUrl + "/chat/completions")
                    .header("Authorization", "Bearer " + apiKey)
                    .header("Content-Type", "application/json")
                    .post(RequestBody.create(jsonBody, MediaType.parse("application/json")))
                    .build();
            
            // 发送请求
            try (Response response = client.newCall(httpRequest).execute()) {
                if (!response.isSuccessful()) {
                    log.error("GLM-4 API call failed: {} - {}", response.code(), response.message());
                    String errorBody = response.body() != null ? response.body().string() : "No error body";
                    log.error("Error body: {}", errorBody);
                    return "Error: GLM-4 API call failed with code " + response.code();
                }
                
                String responseBody = response.body().string();
                log.debug("GLM-4 response: {}", responseBody);
                return parseResponse(responseBody);
            }
            
        } catch (Exception e) {
            log.error("Failed to call GLM-4 API", e);
            return "Error: " + e.getMessage();
        }
    }
    
    /**
     * 构建请求体
     */
    private Map<String, Object> buildRequestBody(LLMRequest request) {
        Map<String, Object> body = new HashMap<>();
        
        // 设置模型
        String modelName = request.getModel() != null ? request.getModel() : model;
        body.put("model", modelName);
        
        // 构建消息列表
        List<Map<String, String>> messages = new ArrayList<>();
        
        // 系统提示
        if (request.getSystemPrompt() != null && !request.getSystemPrompt().isEmpty()) {
            Map<String, String> systemMsg = new HashMap<>();
            systemMsg.put("role", "system");
            systemMsg.put("content", request.getSystemPrompt());
            messages.add(systemMsg);
        }
        
        // 用户提示
        Map<String, String> userMsg = new HashMap<>();
        userMsg.put("role", "user");
        userMsg.put("content", request.getPrompt());
        messages.add(userMsg);
        
        body.put("messages", messages);
        
        // 设置参数
        if (request.getTemperature() != null) {
            body.put("temperature", request.getTemperature());
        }
        
        if (request.getMaxTokens() != null) {
            body.put("max_tokens", request.getMaxTokens());
        }
        
        if (request.getTopP() != null) {
            body.put("top_p", request.getTopP());
        }
        
        if (request.getStopSequences() != null && !request.getStopSequences().isEmpty()) {
            body.put("stop", request.getStopSequences());
        }
        
        // GLM-4 特有参数
        body.put("stream", false); // 不使用流式输出
        
        return body;
    }
    
    /**
     * 解析响应
     */
    private String parseResponse(String responseBody) {
        try {
            @SuppressWarnings("unchecked")
            Map<String, Object> response = objectMapper.readValue(responseBody, Map.class);
            
            @SuppressWarnings("unchecked")
            List<Map<String, Object>> choices = (List<Map<String, Object>>) response.get("choices");
            
            if (choices != null && !choices.isEmpty()) {
                @SuppressWarnings("unchecked")
                Map<String, Object> message = (Map<String, Object>) choices.get(0).get("message");
                if (message != null) {
                    return (String) message.get("content");
                }
            }
            
            return "No response generated";
        } catch (Exception e) {
            log.error("Failed to parse GLM-4 response", e);
            return "Error parsing response: " + e.getMessage();
        }
    }
}

